from ultralytics import YOLO
from swanlab.integration.ultralytics import add_swanlab_callback
import swanlab
from PIL import Image
import time

from conf import operationConfig
from conf import setting
from common import db
from common import llm

class TomotaModel:
    def __init__(self, project, experiment_name, enable_swanlab=False) -> None:
        self.confOperation = operationConfig.OperationConfig()
        self.model = YOLO(f"{setting.DIR_BASE}/yolov8n.pt")
        if enable_swanlab:
            swanlab.init(project=project, experiment_name=experiment_name)
            add_swanlab_callback(self.model)
        self.database = db.Database.get_instance()
        pass
    
    def train(self, epoch, batch, device="cpu"):
        self.model.train(
            data=setting.FILE_PATH["DATA_YAML"],
            epochs=epoch,
            imgsz = 640,
            batch=batch,
            device=device
        )
        self.model.save(self.confOperation.get_mode_path())
        return self

    def val(self, data=setting.FILE_PATH["DATA_YAML"]):
        metrics = self.model.val(data=data)
        print("\n Model Evaluation Metrics:")
        print(f"- Precision: {metrics.box.mp:.3f}")
        print(f"- Recall: {metrics.box.mr:.3f}")
        print(f"- mAP@50: {metrics.box.map50:.3f}")
        print(f"- mAP@50-95: {metrics.box.map:.3f}")  
        print("\n Training & Evaluation Completed! 🚀")

    def predict(self, image, conf_threshold, iou_threshold):
        results = self.model.predict(
            source=image,
            conf = conf_threshold,
            iou = iou_threshold,
            imgsz = 640
        )
        class_results = []
        im = None
        confidence = 0
        for r in results:
            im_array = r.plot()
            im = Image.fromarray(im_array[..., ::-1])
            if r.boxes:
                for box in r.boxes:
                    class_id = int(box.cls)
                    class_name = self.model.names[class_id]
                    confidence = float(box.conf)
                    class_results.append(f'{setting.CHINESE_CLASSNAME[class_name]}({confidence:.2f})')
        unique_results = ", ".join(sorted(set(class_results), key=class_results.index)) if class_results else "The target was not detected"
        row = db.Detection(result=unique_results, confidence=confidence, timestamp=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
        self.database.insert([row])
        return im, unique_results

    @classmethod
    def load_model(cls):
        model = TomotaModel("","")
        model.model = YOLO(operationConfig.OperationConfig().get_mode_path())
        return model